#include <cuda.h>
#if CUDA_VERSION >= 10010


#include <cassert>
#include <cstring>
#include <vector>

#include "NvInfer.h"
#include "kernel/geluKernel.h"
#include "kernel/serialize.h"

static const char* GELE_PLUGIN_VERSION = "1.0";
static const char* GELE_PLUGIN_NAME    = "CTyunGelu";


//////////////////////////////////////////////////////////////////////////
// GeLuPlugin
GeLuPlugin::GeLuPlugin(const std::string name, const nvinfer1::DataType dataType) :
    _pluginName (name), _dataType (dataType) { }

GeLuPlugin::GeLuPlugin(const std::string name, const void* serialData, size_t serialLength) :
    _pluginName (name) {
    deserialize_value(&serialData, &serialLength, &_dataType); 
    deserialize_value(&serialData, &serialLength, &_ld); 
}

nvinfer1::IPluginV2DynamicExt* GeLuPlugin::clone() const noexcept {
    auto plugin = new GeLuPlugin(*this);
    plugin->setPluginNamespace(_namespace.c_str());
    return plugin;
}

nvinfer1::DimsExprs GeLuPlugin::getOutputDimensions(int32_t outputIndex, 
                                        nvinfer1::DimsExprs const* inputs, 
                                        int32_t nbInputs, 
                                        nvinfer1::IExprBuilder& exprBuilder) noexcept {
    return inputs[0];
}


bool GeLuPlugin::supportsFormatCombination(int32_t pos, 
                            nvinfer1::PluginTensorDesc const* inOut, 
                            int32_t nbInputs, 
                            int32_t nbOutputs) noexcept {
    const nvinfer1::PluginTensorDesc& io = inOut[pos];
    return (io.type == _dataType) && (io.format == nvinfer1::TensorFormat::kLINEAR);
}

void GeLuPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, 
                    int32_t nbInputs,
                    nvinfer1::DynamicPluginTensorDesc const* out, 
                    int32_t nbOutputs) noexcept {
    assert(_dataType ==   in[0].desc.type);
    assert(_dataType ==  out[0].desc.type);
}

size_t GeLuPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, 
                        int32_t nbInputs,
                        nvinfer1::PluginTensorDesc const* outputs, 
                        int32_t nbOutputs) const noexcept {
    return 0;
}

int32_t GeLuPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, 
                nvinfer1::PluginTensorDesc const* outputDesc,
                void const* const* inputs, 
                void* const* outputs, 
                void* workspace, 
                cudaStream_t stream) noexcept {
    const int vol = volume(inputDesc[0].dims);
    int status = -1;
    status = geluKernelLauncher(stream,
                        reinterpret_cast<const float*>(inputs[0]), 
                        reinterpret_cast<float*>(outputs[0]),
                        vol
                    );
    return status;
}

nvinfer1::DataType GeLuPlugin::getOutputDataType(int32_t index, 
                                    nvinfer1::DataType const* inputTypes, 
                                    int32_t nbInputs)  const noexcept {
    return inputTypes[0];
}

// IPluginV2 Methods
char const* GeLuPlugin::getPluginType()    const noexcept  { return GELE_PLUGIN_NAME; }
char const* GeLuPlugin::getPluginVersion() const noexcept  { return GELE_PLUGIN_VERSION; }
int32_t     GeLuPlugin::getNbOutputs()     const noexcept  { return 1;}

int32_t     GeLuPlugin::initialize() noexcept { return 0;}
void        GeLuPlugin::terminate()  noexcept { }
void        GeLuPlugin::destroy()    noexcept { delete this; }

size_t      GeLuPlugin::getSerializationSize() const noexcept  {
    return serialized_size<nvinfer1::DataType>(_dataType) + serialized_size<size_t>(_ld);
}

void        GeLuPlugin::serialize(void* buffer) const noexcept  {
    serialize_value(&buffer, _dataType);
    serialize_value(&buffer, _ld);
}

void        GeLuPlugin::setPluginNamespace(const char* pluginNamespace) noexcept {
    _namespace = pluginNamespace;
}

char const* GeLuPlugin::getPluginNamespace() const noexcept{
    return _namespace.c_str();
}

//////////////////////////////////////////////////////////////////////////
// GeLuPluginCreater
nvinfer1::PluginFieldCollection GeLuPluginCreater::_fc{};
std::vector<nvinfer1::PluginField> GeLuPluginCreater::_pluginAttributes{};
REGISTER_TENSORRT_PLUGIN(GeLuPluginCreater);


GeLuPluginCreater::GeLuPluginCreater() {
    _fc.nbFields = _pluginAttributes.size();
    _fc.fields   = _pluginAttributes.data();
}

const char* GeLuPluginCreater::getPluginName()    const noexcept {
    return GELE_PLUGIN_NAME;
}

const char* GeLuPluginCreater::getPluginVersion() const noexcept {
    return GELE_PLUGIN_VERSION;
}

const nvinfer1::PluginFieldCollection* GeLuPluginCreater::getFieldNames() noexcept {
    return &_fc;
}

nvinfer1::IPluginV2DynamicExt* GeLuPluginCreater::createPlugin(const char* name, 
    const nvinfer1::PluginFieldCollection* fc) noexcept {
    int type_id = -1;
    for (int i = 0; i < fc->nbFields; ++i) {
        std::string field_name(fc->fields[i].name);
        if (field_name == "type_id") {
            type_id = *static_cast<const int*>(fc->fields[i].data);
            break;
        }
    }
    if (type_id < 0 || type_id > 3) {
        std::cout <<  "Invalid Type ID." << std::endl;
        return nullptr;
    }
    return new GeLuPlugin(name, static_cast<nvinfer1::DataType>(type_id));
}

nvinfer1::IPluginV2DynamicExt* GeLuPluginCreater::deserializePlugin(const char* name, 
    const void* serialData, size_t serialLength) noexcept {
    return new GeLuPlugin(name, serialData, serialLength);
}

void  GeLuPluginCreater::setPluginNamespace(const char* pluginNamespace) noexcept {
    _namespace = pluginNamespace;
}

const char* GeLuPluginCreater::getPluginNamespace() const noexcept {
    return _namespace.c_str();
}



#endif // CUDA_VERSION >= 100